/**
 * Project Name: socket-test
 * Package Name: com.meatball
 * File Name: EchoHandler
 * Date: 2021/1/19 22:43
 * User: 張翔宇
 * Copyright (c) 2021, zhang.xiangyu@foxmail.com All Rights Reserved.
 */
package com.meatball;

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.socket.HandshakeInfo;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.InetSocketAddress;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 
 * @title: EchoHandler.java
 * @author: 張翔宇
 * @date: 2021/1/19 22:43 
 * @version: V3.0
 */
@Slf4j
@Component
@AllArgsConstructor
@WebSocketMapping("/echo")
public class EchoHandler implements WebSocketHandler {
    private final ConcurrentHashMap<String, WebSocketSender> senderMap;

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        HandshakeInfo handshakeInfo = session.getHandshakeInfo();
        InetSocketAddress address = handshakeInfo.getRemoteAddress();
        Map<String, String> queryMap = getQueryMap(handshakeInfo.getUri().getQuery());
        String userId = queryMap.getOrDefault("userId", "defaultId");

        // 对于输出流：服务器每秒向客户端发送一个数字；
        Mono<Void> input = session.receive()
                .doOnSubscribe(conn -> log.info("建立链接：{}，用户IP：{}，用户：{}", session.getId(),
                        address.getHostName(), userId))
                .map(WebSocketMessage::getPayloadAsText).map(msg -> userId + ": " + msg)
                .then();

        // 对于输入流：每当收到客户端消息时，就打印到标准输出
        Mono<Void> output = session.send(Flux.create(sink ->
                senderMap.put(userId, new WebSocketSender(session, sink)))).log();

         // Mono.zip() 会将多个 Mono 合并为一个新的 Mono，
         // 任何一个 Mono 产生 error 或 complete 都会导致合并后的 Mono
         // 也随之产生 error 或 complete，此时其它的 Mono 则会被执行取消操作。
        return Mono.zip(input, output).then();
    }

    private Map<String, String> getQueryMap(String queryStr) {
        Map<String, String> queryMap = new HashMap<>(1);
        if (StringUtils.hasLength(queryStr)) {
            String[] queryParam = queryStr.split("&");
            Arrays.stream(queryParam).forEach(s -> {
                String[] kv = s.split("=", 2);
                String value = kv.length == 2 ? kv[1] : "";
                queryMap.put(kv[0], value);
            });
        }
        return queryMap;
    }
}
